{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ISubpr_SSsiM"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "3jTMb1dySr3V"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "6DWfyNThSziV"
      },
      "source": [
        "# Better performance with tf.function\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/function\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/function.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/function.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/function.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "J122XQYG7W6w"
      },
      "source": [
        "In TensorFlow 2, eager execution is turned on by default. The user interface is intuitive and flexible (running one-off operations is much easier\n",
        "and faster), but this can come at the expense of performance and deployability.\n",
        "\n",
        "To get performant and portable models, use `tf.function` to make graphs out of your programs. However there are pitfalls to be wary of - `tf.function` is not a magical make-it-faster bullet!\n",
        "\n",
        "This document will help you conceptualize what `tf.function` is doing under the hood, so that you can master its use.\n",
        "\n",
        "The main takeaways and recommendations are:\n",
        "\n",
        "- Debug in Eager mode, then decorate with `@tf.function`.\n",
        "- Don't rely on Python side effects like object mutation or list appends.\n",
        "- tf.function works best with TensorFlow ops; NumPy and Python calls are converted to constants.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SjvqpgepHJPd"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "otIdN1TS8N7S"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "I0xDjO4SHLUD"
      },
      "source": [
        "Define a helper function to demonstrate the kinds of errors you might encounter:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "D25apou9IOXa"
      },
      "outputs": [],
      "source": [
        "import traceback\n",
        "import contextlib\n",
        "\n",
        "# Some helper code to demonstrate the kinds of errors you might encounter.\n",
        "@contextlib.contextmanager\n",
        "def assert_raises(error_class):\n",
        "  try:\n",
        "    yield\n",
        "  except error_class as e:\n",
        "    print('Caught expected exception \\n  {}:'.format(error_class))\n",
        "    traceback.print_exc(limit=2)\n",
        "  except Exception as e:\n",
        "    raise e\n",
        "  else:\n",
        "    raise Exception('Expected {} to be raised but no error was raised!'.format(\n",
        "        error_class))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "WPSfepzTHThq"
      },
      "source": [
        "## Basics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "CNwYTIJ8r56W"
      },
      "source": [
        "A `tf.function` you define is just like a core TensorFlow operation: You can execute it eagerly; you can compute gradients; and so on."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "SbtT1-Wm70F2"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def add(a, b):\n",
        "  return a + b\n",
        "\n",
        "add(tf.ones([2, 2]), tf.ones([2, 2]))  #  [[2., 2.], [2., 2.]]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "uP-zUelB8DbX"
      },
      "outputs": [],
      "source": [
        "v = tf.Variable(1.0)\n",
        "with tf.GradientTape() as tape:\n",
        "  result = add(v, 1.0)\n",
        "tape.gradient(result, v)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ocWZvqrmHnmX"
      },
      "source": [
        "You can use functions inside functions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "l5qRjdbBVdU6"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def dense_layer(x, w, b):\n",
        "  return add(tf.matmul(x, w), b)\n",
        "\n",
        "dense_layer(tf.ones([3, 2]), tf.ones([2, 2]), tf.ones([2]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "piBhz7gYsHqU"
      },
      "source": [
        "Functions can be faster than eager code, especially for graphs with many small ops. But for graphs with a few expensive ops (like convolutions), you may not see much speedup.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "zuXt4wRysI03"
      },
      "outputs": [],
      "source": [
        "import timeit\n",
        "conv_layer = tf.keras.layers.Conv2D(100, 3)\n",
        "\n",
        "@tf.function\n",
        "def conv_fn(image):\n",
        "  return conv_layer(image)\n",
        "\n",
        "image = tf.zeros([1, 200, 200, 100])\n",
        "# warm up\n",
        "conv_layer(image); conv_fn(image)\n",
        "print(\"Eager conv:\", timeit.timeit(lambda: conv_layer(image), number=10))\n",
        "print(\"Function conv:\", timeit.timeit(lambda: conv_fn(image), number=10))\n",
        "print(\"Note how there's not much difference in performance for convolutions\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "aIKkgr6qdtp4"
      },
      "source": [
        "## Debugging\n",
        "\n",
        "In general, debugging code is easier in Eager mode than inside a `tf.function`. You should ensure that your code executes error-free in Eager mode before decorating with `tf.function`. To assist in the debugging process, you can call `tf.config.run_functions_eagerly(True)` to globally disable and reenable `tf.function`.\n",
        "\n",
        "When tracking down issues that only appear within `tf.function`, here are some tips:\n",
        "- Plain old Python `print` calls only execute during tracing, helping you track down when your functions get (re)traced.\n",
        "- `tf.print` calls will execute every time, and can help you track down intermediate values during execution.\n",
        "- `tf.debugging.enable_check_numerics` is an easy way to track down where NaNs and Inf are created.\n",
        "- `pdb` can help you understand what's going on during tracing. (Caveat: PDB will drop you into AutoGraph-transformed source code.)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "uZ4Do2AV80cO"
      },
      "source": [
        "## Tracing and polymorphism\n",
        "\n",
        "Python's dynamic typing means that you can call functions with a variety of argument types, and Python will do something different in each scenario.\n",
        "\n",
        "On the other hand, TensorFlow graphs require static dtypes and shape dimensions. `tf.function` bridges this gap by retracing the function when necessary to generate the correct graphs. Most of the subtlety of `tf.function` usage stems from this retracing behavior.\n",
        "\n",
        "You can call a function with arguments of different types to see what is happening."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "kojmJrgq8U9v"
      },
      "outputs": [],
      "source": [
        "# Functions are polymorphic\n",
        "\n",
        "@tf.function\n",
        "def double(a):\n",
        "  print(\"Tracing with\", a)\n",
        "  return a + a\n",
        "\n",
        "print(double(tf.constant(1)))\n",
        "print()\n",
        "print(double(tf.constant(1.1)))\n",
        "print()\n",
        "print(double(tf.constant(\"a\")))\n",
        "print()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4pJqkDR_Q2wz"
      },
      "source": [
        "To control the tracing behavior, you can use the following techniques:\n",
        "\n",
        "Create a new `tf.function`. Separate `tf.function` objects are guaranteed not to share traces."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "uHp4ousu4DdN"
      },
      "outputs": [],
      "source": [
        "def f():\n",
        "  print('Tracing!')\n",
        "  tf.print('Executing')\n",
        "\n",
        "tf.function(f)()\n",
        "tf.function(f)()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "96IxS2WR37fF"
      },
      "source": [
        "Use `get_concrete_function` method to get a specific trace.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mHg2CGtPQ3Hz"
      },
      "outputs": [],
      "source": [
        "print(\"Obtaining concrete trace\")\n",
        "double_strings = double.get_concrete_function(tf.TensorSpec(shape=None, dtype=tf.string))\n",
        "print(\"Executing traced function\")\n",
        "print(double_strings(tf.constant(\"a\")))\n",
        "print(double_strings(a=tf.constant(\"b\")))\n",
        "print(\"Using a concrete trace with incompatible types will throw an error\")\n",
        "with assert_raises(tf.errors.InvalidArgumentError):\n",
        "  double_strings(tf.constant(1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "EUtycWJa34TT"
      },
      "source": [
        "Specify `input_signature` in `tf.function` to limit tracing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "_BDMIRmu1RGB"
      },
      "outputs": [],
      "source": [
        "@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))\n",
        "def next_collatz(x):\n",
        "  print(\"Tracing with\", x)\n",
        "  return tf.where(x % 2 == 0, x // 2, 3 * x + 1)\n",
        "\n",
        "print(next_collatz(tf.constant([1, 2])))\n",
        "# We specified a 1-D tensor in the input signature, so this should fail.\n",
        "with assert_raises(ValueError):\n",
        "  next_collatz(tf.constant([[1, 2], [3, 4]]))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Es0WZkLIUSdu"
      },
      "source": [
        "## When to retrace?\n",
        "\n",
        "A polymorphic `tf.function` keeps a cache of concrete functions generated by tracing. The cache keys are effectively tuples of keys generated from the function args and kwargs. The key generated for a `tf.Tensor` argument is its number of dimensions and type. The key generated for a Python primitive is its value. For all other Python types, the keys are based on the object `id()` so that methods are traced independently for each instance of a class. In the future, TensorFlow may add more sophisticated cachi\n",
        "ng for Python objects that can be safely converted to tensors.\n",
        "\n",
        "See [Concrete functions](../../guide/concrete_function.ipynb)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "AY5oiQN0XIyA"
      },
      "source": [
        "## Python or Tensor args?\n",
        "\n",
        "Often, Python arguments are used to control hyperparameters and graph constructions - for example, `num_layers=10` or `training=True` or `nonlinearity='relu'`. So if the Python argument changes, it makes sense that you'd have to retrace the graph.\n",
        "\n",
        "However, it's possible that a Python argument is not being used to control graph construction. In these cases, a change in the Python value can trigger needless retracing. Take, for example, this training loop, which AutoGraph will dynamically unroll. Despite the multiple traces, the generated graph is actually identical, so this is a bit inefficient."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "uydzR5JYUU8H"
      },
      "outputs": [],
      "source": [
        "def train_one_step():\n",
        "  pass\n",
        "\n",
        "@tf.function\n",
        "def train(num_steps):\n",
        "  print(\"Tracing with num_steps = {}\".format(num_steps))\n",
        "  for _ in tf.range(num_steps):\n",
        "    train_one_step()\n",
        "\n",
        "train(num_steps=10)\n",
        "train(num_steps=20)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "f6pjnylLUW8P"
      },
      "source": [
        "The simple workaround here is to cast your arguments to Tensors if they do not affect the shape of the generated graph."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "TmL8T-w3UYes"
      },
      "outputs": [],
      "source": [
        "train(num_steps=tf.constant(10))\n",
        "train(num_steps=tf.constant(20))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "129-iRsPS-gY"
      },
      "source": [
        "## Side effects in `tf.function`\n",
        "\n",
        "In general, Python side effects (like printing or mutating objects) only happen during tracing. So how can you reliably trigger side effects from `tf.function`?\n",
        "\n",
        "The general rule of thumb is to only use Python side effects to debug your traces. Otherwise, TensorFlow ops like `tf.Variable.assign`, `tf.print`, and `tf.summary` are the best way to ensure your code will be traced and executed by the TensorFlow runtime with each call. In general using a functional style will yield the best results. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "w2sACuZ9TTRk"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def f(x):\n",
        "  print(\"Traced with\", x)\n",
        "  tf.print(\"Executed with\", x)\n",
        "\n",
        "f(1)\n",
        "f(1)\n",
        "f(2)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "e1I0dPiqTV8H"
      },
      "source": [
        "If you would like to execute Python code during each invocation of a `tf.function`, `tf.py_function` is an exit hatch. The drawback of `tf.py_function` is that it's not portable or particularly performant, nor does it work well in distributed (multi-GPU, TPU) setups. Also, since `tf.py_function` has to be wired into the graph for differentiability, it casts all inputs/outputs to tensors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "7aJD--9qTWmg"
      },
      "outputs": [],
      "source": [
        "external_list = []\n",
        "\n",
        "def side_effect(x):\n",
        "  print('Python side effect')\n",
        "  external_list.append(x)\n",
        "\n",
        "@tf.function\n",
        "def f(x):\n",
        "  tf.py_function(side_effect, inp=[x], Tout=[])\n",
        "\n",
        "f(1)\n",
        "f(1)\n",
        "f(1)\n",
        "assert len(external_list) == 3\n",
        "# .numpy() call required because py_function casts 1 to tf.constant(1)\n",
        "assert external_list[0].numpy() == 1\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "msTmv-oyUNaf"
      },
      "source": [
        "## Beware of Python state\n",
        "\n",
        "Many Python features, such as generators and iterators, rely on the Python runtime to keep track of state. In general, while these constructs work as expected in Eager mode, many unexpected things can happen inside a `tf.function` due to tracing behavior.\n",
        "\n",
        "To give one example, advancing iterator state is a Python side effect and therefore only happens during tracing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "FNPD4unZUedH"
      },
      "outputs": [],
      "source": [
        "external_var = tf.Variable(0)\n",
        "@tf.function\n",
        "def buggy_consume_next(iterator):\n",
        "  external_var.assign_add(next(iterator))\n",
        "  tf.print(\"Value of external_var:\", external_var)\n",
        "\n",
        "iterator = iter([0, 1, 2, 3])\n",
        "buggy_consume_next(iterator)\n",
        "# This reuses the first value from the iterator, rather than consuming the next value.\n",
        "buggy_consume_next(iterator)\n",
        "buggy_consume_next(iterator)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "lPr_6mK_AQWL"
      },
      "source": [
        "## Variables\n",
        "\n",
        "We can use the same idea of leveraging the intended execution order of the code to make variable creation and utilization very easy in `tf.function`. There is one very important caveat, though, which is that with variables it's possible to write code which behaves differently in eager mode and graph mode.\n",
        "\n",
        "Specifically, this will happen when you create a new Variable with each call. Due to tracing semantics, `tf.function` will reuse the same variable each call, but eager mode will create a new variable with each call. To guard against this mistake, `tf.function` will raise an error if it detects dangerous variable creation behavior."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Tx0Vvnb_9OB-"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def f(x):\n",
        "  v = tf.Variable(1.0)\n",
        "  v.assign_add(x)\n",
        "  return v\n",
        "\n",
        "with assert_raises(ValueError):\n",
        "  f(1.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ot8o7PeFIINo"
      },
      "source": [
        "Non-ambiguous code is ok, though."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "DKzNjVg8h4ao"
      },
      "outputs": [],
      "source": [
        "v = tf.Variable(1.0)\n",
        "\n",
        "@tf.function\n",
        "def f(x):\n",
        "  return v.assign_add(x)\n",
        "\n",
        "print(f(1.0))  # 2.0\n",
        "print(f(2.0))  # 4.0\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "KYm6-5GCILXQ"
      },
      "source": [
        "You can also create variables inside a tf.function as long as we can prove\n",
        "that those variables are created only the first time the function is executed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "HQrG5_kOiKl_"
      },
      "outputs": [],
      "source": [
        "class C:\n",
        "  pass\n",
        "\n",
        "obj = C()\n",
        "obj.v = None\n",
        "\n",
        "@tf.function\n",
        "def g(x):\n",
        "  if obj.v is None:\n",
        "    obj.v = tf.Variable(1.0)\n",
        "  return obj.v.assign_add(x)\n",
        "\n",
        "print(g(1.0))  # 2.0\n",
        "print(g(2.0))  # 4.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SDT4oKJaIim8"
      },
      "source": [
        "Variable initializers can depend on function arguments and on values of other\n",
        "variables. We can figure out the right initialization order using the same\n",
        "method we use to generate control dependencies."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "_IOVc1eujMH2"
      },
      "outputs": [],
      "source": [
        "state = []\n",
        "@tf.function\n",
        "def fn(x):\n",
        "  if not state:\n",
        "    state.append(tf.Variable(2.0 * x))\n",
        "    state.append(tf.Variable(state[0] * 3.0))\n",
        "  return state[0] * x * state[1]\n",
        "\n",
        "print(fn(tf.constant(1.0)))\n",
        "print(fn(tf.constant(3.0)))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "5f05Vr_YBUCz"
      },
      "source": [
        "## AutoGraph Transformations\n",
        "\n",
        "AutoGraph is a library that is on by default in `tf.function`, and transforms a subset of Python Eager code into graph-compatible TensorFlow ops. This includes control flow like `if`, `for`, `while`.\n",
        "\n",
        "TensorFlow ops like `tf.cond` and `tf.while_loop` continue to work, but control flow is often easier to write and understand when written in Python."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "yCQTtTPTW3WF"
      },
      "outputs": [],
      "source": [
        "# Simple loop\n",
        "\n",
        "@tf.function\n",
        "def f(x):\n",
        "  while tf.reduce_sum(x) > 1:\n",
        "    tf.print(x)\n",
        "    x = tf.tanh(x)\n",
        "  return x\n",
        "\n",
        "f(tf.random.uniform([5]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "KxwJ8znPI0Cg"
      },
      "source": [
        "If you're curious you can inspect the code autograph generates."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "jlQD1ffRXJhl"
      },
      "outputs": [],
      "source": [
        "print(tf.autograph.to_code(f.python_function))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xgKmkrNTZSyz"
      },
      "source": [
        "### Conditionals\n",
        "\n",
        "AutoGraph will convert some `if <condition>` statements into the equivalent `tf.cond` calls. This substitution is made if `<condition>` is a Tensor. Otherwise, the `if` statement is executed as a Python conditional.\n",
        "\n",
        "A Python conditional executes during tracing, so exactly one branch of the conditional will be added to the graph. Without AutoGraph, this traced graph would be unable to take the alternate branch if there is data-dependent control flow.\n",
        "\n",
        "`tf.cond` traces and adds both branches of the conditional to the graph, dynamically selecting a branch at execution time. Tracing can have unintended side effects; see [AutoGraph tracing effects](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/control_flow.md#effects-of-the-tracing-process) for more."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "BOQl8PMq2Sf3"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def fizzbuzz(n):\n",
        "  for i in tf.range(1, n + 1):\n",
        "    print('Tracing for loop')\n",
        "    if i % 15 == 0:\n",
        "      print('Tracing fizzbuzz branch')\n",
        "      tf.print('fizzbuzz')\n",
        "    elif i % 3 == 0:\n",
        "      print('Tracing fizz branch')\n",
        "      tf.print('fizz')\n",
        "    elif i % 5 == 0:\n",
        "      print('Tracing buzz branch')\n",
        "      tf.print('buzz')\n",
        "    else:\n",
        "      print('Tracing default branch')\n",
        "      tf.print(i)\n",
        "\n",
        "fizzbuzz(tf.constant(5))\n",
        "fizzbuzz(tf.constant(20))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4rBO5AQ15HVC"
      },
      "source": [
        "See the [reference documentation](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/control_flow.md#if-statements) for additional restrictions on AutoGraph-converted if statements."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yho4J0a0ZkQS"
      },
      "source": [
        "### Loops\n",
        "\n",
        "AutoGraph will convert some `for` and `while` statements into the equivalent TensorFlow looping ops, like `tf.while_loop`. If not converted, the `for` or `while` loop is executed as a Python loop.\n",
        "\n",
        "This substitution is made in the following situations:\n",
        "\n",
        "- `for x in y`: if `y` is a Tensor, convert to `tf.while_loop`. In the special case where `y` is a `tf.data.Dataset`, a combination of `tf.data.Dataset` ops are generated.\n",
        "- `while <condition>`: if `<condition>` is a Tensor, convert to `tf.while_loop`.\n",
        "\n",
        "A Python loop executes during tracing, adding additional ops to the `tf.Graph` for every iteration of the loop.\n",
        "\n",
        "A TensorFlow loop traces the body of the loop, and dynamically selects how many iterations to run at execution time.  The loop body only appears once in the generated `tf.Graph`.\n",
        "\n",
        "See the [reference documentation](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/control_flow.md#while-statements) for additional restrictions on AutoGraph-converted `for` and `while` statements."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "sp4rbIdfbM6s"
      },
      "source": [
        "#### Looping over Python data\n",
        "\n",
        "A common pitfall is to loop over Python/Numpy data within a `tf.function`. This loop will execute during the tracing process, adding a copy of your model to the `tf.Graph` for each iteration of the loop.\n",
        "\n",
        "If you want to wrap the entire training loop in `tf.function`, the safest way to do this is to wrap your data as a `tf.data.Dataset` so that AutoGraph will dynamically unroll the training loop."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "WGZ19LspbZ27"
      },
      "outputs": [],
      "source": [
        "def measure_graph_size(f, *args):\n",
        "  g = f.get_concrete_function(*args).graph\n",
        "  print(\"{}({}) contains {} nodes in its graph\".format(\n",
        "      f.__name__, ', '.join(map(str, args)), len(g.as_graph_def().node)))\n",
        "\n",
        "@tf.function\n",
        "def train(dataset):\n",
        "  loss = tf.constant(0)\n",
        "  for x, y in dataset:\n",
        "    loss += tf.abs(y - x) # Some dummy computation.\n",
        "  return loss\n",
        "\n",
        "small_data = [(1, 1)] * 3\n",
        "big_data = [(1, 1)] * 10\n",
        "measure_graph_size(train, small_data)\n",
        "measure_graph_size(train, big_data)\n",
        "\n",
        "measure_graph_size(train, tf.data.Dataset.from_generator(\n",
        "    lambda: small_data, (tf.int32, tf.int32)))\n",
        "measure_graph_size(train, tf.data.Dataset.from_generator(\n",
        "    lambda: big_data, (tf.int32, tf.int32)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JeD2U-yrbfVb"
      },
      "source": [
        "When wrapping Python/Numpy data in a Dataset, be mindful of `tf.data.Dataset.from_generator` versus ` tf.data.Dataset.from_tensors`. The former will keep the data in Python and fetch it via `tf.py_function` which can have performance implications, whereas the latter will bundle a copy of the data as one large `tf.constant()` node in the graph, which can have memory implications.\n",
        "\n",
        "Reading data from files via TFRecordDataset/CsvDataset/etc. is the most effective way to consume data, as then TensorFlow itself can manage the asynchronous loading and prefetching of data, without having to involve Python. To learn more, see the [tf.data guide](../../guide/data)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hyksHW9TCukR"
      },
      "source": [
        "#### Accumulating values in a loop\n",
        "\n",
        "A common pattern is to accumulate intermediate values from a loop. Normally, this is accomplished by appending to a Python list or adding entries to a Python dictionary. However, as these are Python side effects, they will not work as expected in a dynamically unrolled loop. Use `tf.TensorArray` to accumulate results from a dynamically unrolled loop."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "HJ3Vb3dXfefN"
      },
      "outputs": [],
      "source": [
        "batch_size = 2\n",
        "seq_len = 3\n",
        "feature_size = 4\n",
        "\n",
        "def rnn_step(inp, state):\n",
        "  return inp + state\n",
        "\n",
        "@tf.function\n",
        "def dynamic_rnn(rnn_step, input_data, initial_state):\n",
        "  # [batch, time, features] -> [time, batch, features]\n",
        "  input_data = tf.transpose(input_data, [1, 0, 2])\n",
        "  max_seq_len = input_data.shape[0]\n",
        "\n",
        "  states = tf.TensorArray(tf.float32, size=max_seq_len)\n",
        "  state = initial_state\n",
        "  for i in tf.range(max_seq_len):\n",
        "    state = rnn_step(input_data[i], state)\n",
        "    states = states.write(i, state)\n",
        "  return tf.transpose(states.stack(), [1, 0, 2])\n",
        "  \n",
        "dynamic_rnn(rnn_step,\n",
        "            tf.random.uniform([batch_size, seq_len, feature_size]),\n",
        "            tf.zeros([batch_size, feature_size]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "IKyrEY5GVX3M"
      },
      "source": [
        "## Further reading\n",
        "\n",
        "To learn more about graph optimizations that are performed after tracing a `tf.function`, see the [Grappler guide](../../guide/graph_optimization). To learn how to optimize your data pipeline and profile your model, see the [Profiler guide](../../guide/profiler.md)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "function.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
